
# 使用 langchain 解析文档文件（doc, docx, excel-xls\xlsx）

import os
import sys
import tempfile
import subprocess
import shutil
import logging
from pathlib import Path
from typing import Optional
from io import BytesIO
from fastapi import UploadFile
from langchain_community.document_loaders import UnstructuredWordDocumentLoader
import pandas as pd
from openpyxl import Workbook, load_workbook
import docx

logger = logging.getLogger(__name__)

# 添加项目根目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
sys.path.insert(0, project_root)


class RAGFlowExcelParser:
    """
    RAGFlow Excel 解析器
    支持 xlsx, xls, csv 格式的 Excel 文件解析
    """
    
    @staticmethod
    def explode_merged_cells(ws):
        """
        处理合并单元格，将合并区域的值填充到所有单元格中
        """
        for rng in list(ws.merged_cells.ranges):  # 复制，以免迭代中修改
            # 获取左上角单元格的值
            tl_cell = ws.cell(rng.min_row, rng.min_col)
            tl_value = tl_cell.value

            # 先解除合并，否则无法写值
            ws.unmerge_cells(str(rng))

            # 将左上角的值填充到合并区域内的所有单元格
            for r in range(rng.min_row, rng.max_row + 1):
                for c in range(rng.min_col, rng.max_col + 1):
                    ws.cell(r, c).value = tl_value

    @staticmethod
    def _load_excel_to_workbook(file_like_object):
        """加载 Excel 文件到 Workbook 对象"""
        if isinstance(file_like_object, bytes):
            file_like_object = BytesIO(file_like_object)

        # 读取前4字节判断文件类型
        file_like_object.seek(0)
        file_head = file_like_object.read(4)
        file_like_object.seek(0)

        # 如果不是 Excel 格式，尝试作为 CSV 处理
        if not (file_head.startswith(b'PK\x03\x04') or file_head.startswith(b'\xD0\xCF\x11\xE0')):
            logger.info("Not an Excel file, converting CSV to Excel Workbook")
            try:
                file_like_object.seek(0)
                df = pd.read_csv(file_like_object)
                return RAGFlowExcelParser._dataframe_to_workbook(df)  # df转成Workbook好处理
            except Exception as e_csv:
                raise Exception(f"Failed to parse CSV and convert to Excel Workbook: {e_csv}")

        # 尝试使用 openpyxl 加载
        try:
            wb = load_workbook(file_like_object)
            # 处理所有工作表的合并单元格
            for ws in wb.worksheets:
                RAGFlowExcelParser.explode_merged_cells(ws)

            return wb
        except Exception as e:
            logger.info(f"openpyxl load error: {e}, try pandas instead")
            try:
                file_like_object.seek(0)
                df = pd.read_excel(file_like_object)
                return RAGFlowExcelParser._dataframe_to_workbook(df)
            except Exception as e_pandas:
                raise Exception(f"pandas.read_excel error: {e_pandas}, original openpyxl error: {e}")

    @staticmethod
    def _dataframe_to_workbook(df):
        """将 DataFrame 转换为 Workbook 对象"""
        wb = Workbook()
        ws = wb.active
        ws.title = "Data"

        # 写入列标题
        for col_num, column_name in enumerate(df.columns, 1):
            ws.cell(row=1, column=col_num, value=column_name)

        # 写入数据行
        for row_num, row in enumerate(df.values, 2):
            for col_num, value in enumerate(row, 1):
                ws.cell(row=row_num, column=col_num, value=value)

        return wb

    def __call__(self, fnm):
        """解析 Excel 文件并返回文本列表"""
        file_like_object = BytesIO(fnm) if not isinstance(fnm, str) else fnm
        wb = RAGFlowExcelParser._load_excel_to_workbook(file_like_object)  # 返回拆除合并后的 Workbook 对象

        res = []
        for sheetname in wb.sheetnames:
            ws = wb[sheetname]
            rows = list(ws.rows)
            if not rows:
                continue

            # 获取表头
            ti = list(rows[0])

            # 处理数据行
            for r in list(rows[1:]):
                unique_items = []
                seen = set()
                for i, c in enumerate(r):
                    if not c.value:
                        continue
                    field_name = str(ti[i].value) if i < len(ti) else ""
                    field_value = str(c.value)
                    for item in (field_name, field_value):
                        if item and item not in seen:  # 去重并且保留顺序
                            unique_items.append(item)
                            seen.add(item)
                # 用'|'拼接所有去重且有序的field_name和field_value
                line = '|'.join(unique_items)
                if sheetname.lower().find("sheet") < 0:
                    line += " ——" + sheetname  # 最后加上工作表名
                if line.strip():
                    res.append(line)
        return res


def convert_doc_to_docx(doc_path: str, output_dir: Optional[str] = None) -> str:
    """
    将 .doc 文件转换为 .docx 文件（Linux 使用）
    
    Args:
        doc_path: .doc 文件路径
        output_dir: 输出目录，如果为 None 则使用 doc_path 所在目录
        
    Returns:
        str: 转换后的 .docx 文件路径
        
    Raises:
        RuntimeError: 如果转换失败
    """
    if output_dir is None:
        output_dir = os.path.dirname(doc_path)
    
    output_path = os.path.join(output_dir, os.path.basename(doc_path).replace('.doc', '.docx'))
    
    # Linux: 使用 LibreOffice 命令行工具
    soffice = shutil.which('soffice')
    if soffice:
        try:
            cmd = [
                soffice,
                '--headless',
                '--convert-to', 'docx',
                '--outdir', output_dir,
                doc_path
            ]
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
            if result.returncode == 0 and os.path.exists(output_path):
                return output_path
        except Exception as e:
            print(f"Warning: LibreOffice conversion failed: {e}")
    
    # 如果转换失败
    raise RuntimeError(
        "Failed to convert .doc to .docx. "
        "Please install LibreOffice Writer: "
        "sudo apt-get install libreoffice-writer (Ubuntu/Debian) "
        "or sudo dnf install libreoffice-writer (CentOS/RHEL)"
    )

async def parse_document_content(file: UploadFile) -> str:
    """
    使用 langchain 解析文档内容（doc, docx, excel）
    
    Args:
        file: FastAPI UploadFile 对象
        
    Returns:
        str: 文档的文本内容
    """
    
    filename = file.filename or ""
    file_extension = os.path.splitext(filename)[1].lower()
    
    # 重置文件指针到开始位置（如果文件已经被读取过）
    try:
        # FastAPI UploadFile 的 file 属性是底层的文件对象
        if hasattr(file, 'file') and hasattr(file.file, 'seek'):
            file.file.seek(0)
        elif hasattr(file, 'seek'):
            file.seek(0)
    except Exception:
        pass  # 如果重置失败，继续尝试读取
    
    # 创建临时文件
    # Windows 和 Linux 都支持 tempfile
    # 使用 tempfile.gettempdir() 确保跨平台兼容
    temp_dir = tempfile.gettempdir()
    temp_file_path = os.path.join(temp_dir, f"parse_{os.urandom(8).hex()}{file_extension}")
    
    try:
        # 读取上传的文件内容并写入临时文件
        content = await file.read()
        with open(temp_file_path, 'wb') as temp_file:
            temp_file.write(content)
            # 确保文件写入磁盘
            temp_file.flush()
            os.fsync(temp_file.fileno())
        
        # Linux 上设置文件权限（确保可读）
        if os.name != 'nt':  # 非 Windows 系统
            os.chmod(temp_file_path, 0o644)
    except Exception as e:
        # 如果创建临时文件失败，清理并抛出异常
        if os.path.exists(temp_file_path):
            try:
                os.remove(temp_file_path)
            except:
                pass
        raise IOError(f"Failed to create temporary file: {e}")
    
    try:
        documents = []
        
        # 根据文件扩展名选择相应的加载器
        try:
            if file_extension == ".docx":
                # .docx 文件使用 python-docx（无需系统依赖）
                try:
                    doc = docx.Document(temp_file_path)
                    text_parts = []
                    for paragraph in doc.paragraphs:
                        if paragraph.text.strip():
                            text_parts.append(paragraph.text)
                    # 提取表格内容
                    for table in doc.tables:
                        for row in table.rows:
                            row_text = " | ".join(cell.text.strip() for cell in row.cells)
                            if row_text.strip():
                                text_parts.append(row_text)
                    return "\n".join(text_parts)
                except Exception as e:
                    # 如果 python-docx 失败，回退到 langchain
                    print(f"Warning: Failed to parse .docx with python-docx: {e}, trying langchain...")
                    loader = UnstructuredWordDocumentLoader(temp_file_path)
                    documents = loader.load()
            elif file_extension == ".doc":
                # .doc 文件处理：Windows 使用 win32com 转换，Linux 使用 LibreOffice 转换
                if os.name == 'nt':  # Windows
                    # Windows: 使用 win32com 将 .doc 转换为 .docx
                    try:
                        import win32com.client as wc
                        import pythoncom
                        
                        filename_without_ext = os.path.splitext(os.path.basename(temp_file_path))[0]
                        docx_path = os.path.join(temp_dir, f"{filename_without_ext}.docx")
                        
                        pythoncom.CoInitialize()
                        try:
                            word = wc.Dispatch("Word.Application")
                            word.Visible = False
                            try:
                                doc = word.Documents.Open(os.path.abspath(temp_file_path))
                                doc.SaveAs(os.path.abspath(docx_path), 12)  # 12 = wdFormatXMLDocument (.docx)
                                doc.Close()
                                word.Quit()
                            except Exception as e:
                                print(f"打开Word文件失败: {e}")
                                word.Quit()
                                raise
                        finally:
                            pythoncom.CoUninitialize()
                        
                        # 使用 python-docx 解析转换后的 .docx 文件
                        doc = docx.Document(docx_path)
                        text_parts = []
                        
                        # 按文档顺序提取内容（段落和表格混合）
                        # 方法1：使用 XML 遍历保持原始顺序
                        try:
                            from docx.oxml.ns import qn
                            
                            for element in doc.element.body:
                                tag = element.tag
                                if tag.endswith('p'):  # 段落
                                    # 提取段落文本（包括所有嵌套的文本节点）
                                    para_text = ''.join(node.text for node in element.iter() if node.text)
                                    para_text = para_text.strip()
                                    if para_text:
                                        text_parts.append(para_text)
                                elif tag.endswith('tbl'):  # 表格
                                    # 提取表格内容，遍历所有行和单元格
                                    for tr in element.iter(qn('w:tr')):  # 表格行
                                        row_cells = []
                                        for tc in tr.iter(qn('w:tc')):  # 表格单元格
                                            # 提取单元格内的所有文本（包括嵌套的段落）
                                            cell_text = ''.join(node.text for node in tc.iter() if node.text)
                                            cell_text = cell_text.strip()
                                            if cell_text:
                                                row_cells.append(cell_text)
                                        if row_cells:
                                            text_parts.append(" | ".join(row_cells))
                        except Exception as e:
                            print(f"Warning: XML parsing failed: {e}, using fallback method...")
                            text_parts = []
                        
                        # 如果上面的方法没有提取到内容，使用备用方法
                        if not text_parts:
                            # 备用方法：分别提取段落和表格
                            for paragraph in doc.paragraphs:
                                para_text = paragraph.text.strip()
                                if para_text:
                                    text_parts.append(para_text)
                            # 提取表格内容（确保所有表格都被提取）
                            for table in doc.tables:
                                for row in table.rows:
                                    row_cells = []
                                    for cell in row.cells:
                                        # 提取单元格内的所有文本（包括段落）
                                        cell_text = cell.text.strip()
                                        if cell_text:
                                            row_cells.append(cell_text)
                                    if row_cells:
                                        text_parts.append(" | ".join(row_cells))
                        # 清理转换后的临时文件
                        try:
                            if os.path.exists(docx_path):
                                os.remove(docx_path)
                        except:
                            pass
                        return "\n".join(text_parts)
                    except ImportError:
                        # 如果 win32com 未安装，回退到 langchain
                        print("Warning: win32com not installed, falling back to langchain...")
                        loader = UnstructuredWordDocumentLoader(temp_file_path)
                        documents = loader.load()
                    except Exception as e:
                        # 如果转换失败，回退到 langchain
                        print(f"Warning: Failed to convert .doc to .docx with win32com: {e}, trying langchain...")
                        loader = UnstructuredWordDocumentLoader(temp_file_path)
                        documents = loader.load()
                else:  # Linux/Unix
                    # Linux: 转换为 .docx 后再解析
                    try:
                        docx_path = convert_doc_to_docx(temp_file_path, temp_dir)
                        # 使用 python-docx 解析转换后的 .docx 文件
                        doc = docx.Document(docx_path)
                        text_parts = []
                        
                        # 按文档顺序提取内容（段落和表格混合）
                        # 方法1：使用 XML 遍历保持原始顺序
                        try:
                            from docx.oxml.ns import qn
                            
                            for element in doc.element.body:
                                tag = element.tag
                                if tag.endswith('p'):  # 段落
                                    # 提取段落文本（包括所有嵌套的文本节点）
                                    para_text = ''.join(node.text for node in element.iter() if node.text)
                                    para_text = para_text.strip()
                                    if para_text:
                                        text_parts.append(para_text)
                                elif tag.endswith('tbl'):  # 表格
                                    # 提取表格内容，遍历所有行和单元格
                                    for tr in element.iter(qn('w:tr')):  # 表格行
                                        row_cells = []
                                        for tc in tr.iter(qn('w:tc')):  # 表格单元格
                                            # 提取单元格内的所有文本（包括嵌套的段落）
                                            cell_text = ''.join(node.text for node in tc.iter() if node.text)
                                            cell_text = cell_text.strip()
                                            if cell_text:
                                                row_cells.append(cell_text)
                                        if row_cells:
                                            text_parts.append(" | ".join(row_cells))
                        except Exception as e:
                            print(f"Warning: XML parsing failed: {e}, using fallback method...")
                            text_parts = []
                        
                        # 如果上面的方法没有提取到内容，使用备用方法
                        if not text_parts:
                            # 备用方法：分别提取段落和表格
                            for paragraph in doc.paragraphs:
                                para_text = paragraph.text.strip()
                                if para_text:
                                    text_parts.append(para_text)
                            # 提取表格内容（确保所有表格都被提取）
                            for table in doc.tables:
                                for row in table.rows:
                                    row_cells = []
                                    for cell in row.cells:
                                        # 提取单元格内的所有文本（包括段落）
                                        cell_text = cell.text.strip()
                                        if cell_text:
                                            row_cells.append(cell_text)
                                    if row_cells:
                                        text_parts.append(" | ".join(row_cells))
                        # 清理转换后的临时文件
                        try:
                            if os.path.exists(docx_path):
                                os.remove(docx_path)
                        except:
                            pass
                        return "\n".join(text_parts)
                    except Exception as e:
                        # 如果转换失败，回退到 langchain（需要 LibreOffice）
                        print(f"Warning: Failed to convert .doc to .docx: {e}, trying langchain...")
                        loader = UnstructuredWordDocumentLoader(temp_file_path)
                        documents = loader.load()
                
            elif file_extension in [".xls", ".xlsx", ".csv"]:
                # 使用 RAGFlowExcelParser 解析 Excel（跨平台兼容，无需系统依赖）
                try:
                    # 读取文件内容
                    with open(temp_file_path, 'rb') as f:
                        file_content = f.read()
                    
                    # 使用 RAGFlowExcelParser 解析
                    excel_parser = RAGFlowExcelParser()
                    parsed_content = excel_parser(file_content)
                    
                    # 将解析结果合并为文本
                    content_text = "\n".join(parsed_content)
                    return content_text
                except Exception as e:
                    raise ValueError(f"Failed to parse Excel file: {e}")
                
            else:
                raise ValueError(f"Unsupported file type: {file_extension}. Supported types: .doc, .docx, .xls, .xlsx, .csv")
        except ImportError as e:
            # 捕获缺少系统依赖的错误
            error_msg = str(e)
            if "docx" in error_msg.lower() or "No module named" in error_msg:
                raise ImportError(
                    f"Missing dependency: {error_msg}. "
                    "On Linux, you may need to install system packages: "
                    "sudo apt-get install libmagic1 poppler-utils tesseract-ocr libreoffice "
                    "(or equivalent for your distribution)"
                )
            raise
        except Exception as e:
            # 其他错误，提供更友好的错误信息
            error_msg = str(e)
            if "libmagic" in error_msg.lower() or "file" in error_msg.lower():
                raise RuntimeError(
                    f"Document parsing failed: {error_msg}. "
                    "On Linux, you may need to install: sudo apt-get install libmagic1"
                )
            raise
        
        # 提取所有文档的文本内容并合并
        if documents:
            content_text = "\n".join(doc.page_content for doc in documents)
            return content_text
        else:
            return ""

    finally:
        # 清理临时文件
        try:
            if os.path.exists(temp_file_path):
                os.remove(temp_file_path)
        except Exception as e:
            print(f"Warning: Failed to delete temporary file {temp_file_path}: {e}")


async def parse_file_by_path(file_path: str) -> str:
    """
    通过文件路径解析文档内容（用于测试或直接文件路径场景）

    Args:
        file_path: 文件路径

    Returns:
        str: 文档的文本内容
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"File not found: {file_path}")

    file_extension = os.path.splitext(file_path)[1].lower()
    documents = []

    # 根据文件扩展名选择相应的加载器
    if file_extension == ".docx":
        # .docx 文件使用 python-docx（无需系统依赖）
        try:
            doc = docx.Document(file_path)
            text_parts = []
            for paragraph in doc.paragraphs:
                if paragraph.text.strip():
                    text_parts.append(paragraph.text)
            # 提取表格内容
            for table in doc.tables:
                for row in table.rows:
                    row_text = " | ".join(cell.text.strip() for cell in row.cells)
                    if row_text.strip():
                        text_parts.append(row_text)
            return "\n".join(text_parts)
        except Exception as e:
            # 如果 python-docx 失败，回退到 langchain
            print(f"Warning: Failed to parse .docx with python-docx: {e}, trying langchain...")
            loader = UnstructuredWordDocumentLoader(file_path)
            documents = loader.load()
    elif file_extension == ".doc":
        # .doc 文件处理：Windows 使用 win32com 转换，Linux 使用 LibreOffice 转换
        if os.name == 'nt':  # Windows
            # Windows: 使用 win32com 将 .doc 转换为 .docx
            try:
                import win32com.client as wc
                import pythoncom
                
                output_dir = os.path.dirname(file_path)
                filename_without_ext = os.path.splitext(os.path.basename(file_path))[0]
                docx_path = os.path.join(output_dir, f"{filename_without_ext}.docx")
                
                pythoncom.CoInitialize()
                try:
                    word = wc.Dispatch("Word.Application")
                    word.Visible = False
                    try:
                        doc = word.Documents.Open(os.path.abspath(file_path))
                        doc.SaveAs(os.path.abspath(docx_path), 12)  # 12 = wdFormatXMLDocument (.docx)
                        doc.Close()
                        word.Quit()
                    except Exception as e:
                        print(f"打开Word文件失败: {e}")
                        word.Quit()
                        raise
                finally:
                    pythoncom.CoUninitialize()
                
                # 使用 python-docx 解析转换后的 .docx 文件
                doc = docx.Document(docx_path)
                text_parts = []
                for paragraph in doc.paragraphs:
                    if paragraph.text.strip():
                        text_parts.append(paragraph.text)
                # 提取表格内容
                for table in doc.tables:
                    for row in table.rows:
                        row_text = " | ".join(cell.text.strip() for cell in row.cells)
                        if row_text.strip():
                            text_parts.append(row_text)
                # 清理转换后的临时文件（如果是在临时目录）
                if os.path.dirname(docx_path) == tempfile.gettempdir():
                    try:
                        if os.path.exists(docx_path):
                            os.remove(docx_path)
                    except:
                        pass
                return "\n".join(text_parts)
            except ImportError:
                # 如果 win32com 未安装，回退到 langchain
                print("Warning: win32com not installed, falling back to langchain...")
                loader = UnstructuredWordDocumentLoader(file_path)
                documents = loader.load()
            except Exception as e:
                # 如果转换失败，回退到 langchain
                print(f"Warning: Failed to convert .doc to .docx with win32com: {e}, trying langchain...")
                loader = UnstructuredWordDocumentLoader(file_path)
                documents = loader.load()
        else:  # Linux/Unix
            # Linux: 转换为 .docx 后再解析
            try:
                output_dir = os.path.dirname(file_path)
                docx_path = convert_doc_to_docx(file_path, output_dir)
                # 使用 python-docx 解析转换后的 .docx 文件
                doc = docx.Document(docx_path)
                text_parts = []
                for paragraph in doc.paragraphs:
                    if paragraph.text.strip():
                        text_parts.append(paragraph.text)
                # 提取表格内容
                for table in doc.tables:
                    for row in table.rows:
                        row_text = " | ".join(cell.text.strip() for cell in row.cells)
                        if row_text.strip():
                            text_parts.append(row_text)
                # 清理转换后的临时文件（如果是在临时目录）
                if os.path.dirname(docx_path) == tempfile.gettempdir():
                    try:
                        if os.path.exists(docx_path):
                            os.remove(docx_path)
                    except:
                        pass
                return "\n".join(text_parts)
            except Exception as e:
                # 如果转换失败，回退到 langchain（需要 LibreOffice）
                print(f"Warning: Failed to convert .doc to .docx: {e}, trying langchain...")
                loader = UnstructuredWordDocumentLoader(file_path)
                documents = loader.load()

    elif file_extension in [".xls", ".xlsx", ".csv"]:
        # 使用 RAGFlowExcelParser 解析 Excel（跨平台兼容，无需系统依赖）
        try:
            # 读取文件内容
            with open(file_path, 'rb') as f:
                file_content = f.read()
            
            # 使用 RAGFlowExcelParser 解析
            excel_parser = RAGFlowExcelParser()
            parsed_content = excel_parser(file_content)
            
            # 将解析结果合并为文本
            content_text = "\n".join(parsed_content)
            return content_text
        except Exception as e:
            raise ValueError(f"Failed to parse Excel file: {e}")

    else:
        raise ValueError(f"Unsupported file type: {file_extension}. Supported types: .doc, .docx, .xls, .xlsx, .csv")
